#include "deepsort/mot_drawer.h"

namespace {
    const static std::vector<std::string> kLabels = { "car", "van", "bus", "others" };
    // const static std::vector<std::string> kLabels = { "person", "bicycle", "car", "motorbike",
    //   "aeroplane","bus", "train", "truck", "boat",
    //   "traffic light", "fire hydrant", "stop sign", "parking meter",
    //   "bench", "bird", "cat", "dog", "horse",
    //   "sheep", "cow", "elephant", "bear", "zebra",
    //   "giraffe", "backpack", "umbrella", "handbag","tie",
    //   "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    //   "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
    //   "tennis racket", "bottle", "wine glass", "cup",
    //   "fork", "knife", "spoon", "bowl", "banana",
    //   "apple", "sandwich", "orange", "broccoli", "carrot",
    //   "hot dog", "pizza", "donut", "cake", "chair",
    //   "sofa", "potted plant", "bed", "dining table", "toilet",
    //   "TV monitor", "laptop", "mouse", "remote", "keyboard",
    //   "cell phone", "microwave", "oven", "toaster", "sink",
    //   "refrigerator", "book", "clock", "vase","scissors",
    //   "teddy bear", "hair drier", "toothbrush" };

    const std::string OUTPUT_DIR = "./output";
    const std::string TARGETS_DIR_NAME = "targets";
    const std::string TRAJ_TXT_NAME = "trajectories.txt";
    const std::string SNAPSHOTS_DIR_NAME = "snapshots";
}
namespace fs = std::experimental::filesystem;

MotDrawer::MotDrawer() {}

MotDrawer::~MotDrawer() {
    record(0);
}

void MotDrawer::DrawTrajectories(cv::Mat& ori_img, vector<BBox>& trackResults, bool isDrawTrajectory) {

    for (auto &t : trackResults) {
        Utils::draw_bbox(ori_img, t.rect, to_string(t.trk_id)+kLabels[t.cls_id], Utils::color_map(t.trk_id));
        if (isDrawTrajectory) {
            draw_trajectory(ori_img, get().at(t.trk_id).trajectories, Utils::color_map(t.trk_id));   
        }
    }
}

void MotDrawer::draw_trajectory(cv::Mat &img, const std::map<int, cv::Rect2f> &traj,
                        const cv::Scalar &color) {
    if (traj.size() < 2) return;

    auto cur = traj.begin()->second;
    auto pt1 = cur.br();
    pt1.x -= cur.width / 2;
    pt1.x *= img.cols;
    pt1.y *= img.rows;

    for (auto it = ++traj.begin(); it != traj.end(); ++it) {
        cur = it->second;
        auto pt2 = cur.br();
        pt2.x -= cur.width / 2;
        pt2.x *= img.cols;
        pt2.y *= img.rows;
        cv::line(img, pt1, pt2, color);
        pt1 = pt2;
    }
}

// 写视频，写txt结果，写snapshot
void MotDrawer::update(const vector<BBox> &trks, uint32_t frame, const cv::Mat &image) {
    for (auto trk : trks) {
        // save normalized boxes
        cv::Rect2f boxf;

        boxf = normalize_rect(trk.rect, image.cols, image.rows);
        auto &t = targets[trk.trk_id];
        t.trajectories.emplace(frame, boxf);
        if ((frame - t.last_snap) > 5) {
            t.snapshots[frame] = image(unnormalize_rect(pad_rect(boxf, padding), image.cols, image.rows)).clone();
            t.last_snap = frame;
        }
    }
    record(20);
}

void MotDrawer::record(uint32_t remain) {
    for (auto&[id, t] : targets) {
        auto dir_path = fs::path(OUTPUT_DIR) / TARGETS_DIR_NAME / to_string(id);
        fs::create_directories(dir_path);

        ofstream fp(dir_path / TRAJ_TXT_NAME, ios::app);
        fp << right << fixed << setprecision(3);
        while (t.trajectories.size() > remain) {
            auto &[frame, box] = *t.trajectories.begin();
            fp << setw(6) << frame
               << setw(6) << box.x
               << setw(6) << box.y
               << setw(6) << box.width
               << setw(6) << box.height
               << setw(6) << endl;
            t.trajectories.erase(t.trajectories.begin());
        }

        dir_path /= SNAPSHOTS_DIR_NAME;
        fs::create_directories(dir_path);
        while (!t.snapshots.empty()) {
            auto &[frame, ss] = *t.snapshots.begin();
            if (t.trajectories.empty() || frame < t.trajectories.begin()->first) {
                cv::imwrite((dir_path / (to_string(frame) + ".jpg")).string(), ss);
                t.snapshots.erase(t.snapshots.begin());
            }
            else {
                break;
            }
        }
    }
}

cv::Rect2f MotDrawer::pad_rect(cv::Rect2f rect, float padding) {
    rect.x = std::max(0.0f, rect.x - rect.width * padding);
    rect.y = std::max(0.0f, rect.y - rect.height * padding);
    rect.width = std::min(1 - rect.x, rect.width * (1 + 2 * padding));
    rect.height = std::min(1 - rect.y, rect.height * (1 + 2 * padding));

    return rect;
}

cv::Rect2f MotDrawer::normalize_rect(cv::Rect rect, float w, float h) {
    cv::Rect2f res;
    res.x = rect.x / w;
    res.y = rect.y / h;
    res.width = rect.width / w;
    res.height = rect.height / h;
    return res;
}

cv::Rect MotDrawer::unnormalize_rect(cv::Rect2f rect, float w, float h) {
    cv::Rect res;
    res.x = rect.x * w;
    res.y = rect.y * h;
    res.width = rect.width * w;
    res.height = rect.height * h;
    return res;
}
